{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Notebook written by [Zhedong Zheng](https://github.com/zhedongzheng)\n",
    "\n",
    "<img src=\"img/charrnn.jpeg\" width=\"400\">"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import tensorflow as tf\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "params = {\n",
    "    'batch_size': 128,\n",
    "    'text_iter_step': 25,\n",
    "    'seq_len': 200,\n",
    "    'hidden_dim': 128,\n",
    "    'n_layers': 2,\n",
    "    'beam_width': 5,\n",
    "    'display_step': 10,\n",
    "    'generate_step': 100,\n",
    "    'clip_norm': 5.0,\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "def parse_text(file_path):\n",
    "    with open(file_path) as f:\n",
    "        text = f.read()\n",
    "    \n",
    "    char2idx = {c: i+3 for i, c in enumerate(set(text))}\n",
    "    char2idx['<pad>'] = 0\n",
    "    char2idx['<start>'] = 1\n",
    "    char2idx['<end>'] = 2\n",
    "    \n",
    "    ints = np.array([char2idx[char] for char in list(text)])\n",
    "    return ints, char2idx\n",
    "\n",
    "def next_batch(ints):\n",
    "    len_win = params['seq_len'] * params['batch_size']\n",
    "    for i in range(0, len(ints)-len_win, params['text_iter_step']):\n",
    "        clip = ints[i: i+len_win]\n",
    "        yield clip.reshape([params['batch_size'], params['seq_len']])\n",
    "        \n",
    "def input_fn(ints):\n",
    "    dataset = tf.data.Dataset.from_generator(\n",
    "        lambda: next_batch(ints), tf.int32, tf.TensorShape([None, params['seq_len']]))\n",
    "    iterator = dataset.make_one_shot_iterator()\n",
    "    return iterator.get_next()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "def start_sent(x):\n",
    "    _x = tf.fill([tf.shape(x)[0], 1], params['char2idx']['<start>'])\n",
    "    return tf.concat([_x, x], 1)\n",
    "\n",
    "def end_sent(x):\n",
    "    _x = tf.fill([tf.shape(x)[0], 1], params['char2idx']['<end>'])\n",
    "    return tf.concat([x, _x], 1)\n",
    "\n",
    "def cell_fn():\n",
    "    return tf.nn.rnn_cell.ResidualWrapper(\n",
    "        tf.nn.rnn_cell.GRUCell(params['hidden_dim'],\n",
    "            kernel_initializer=tf.orthogonal_initializer()))\n",
    "  \n",
    "def multi_cell_fn():\n",
    "    return tf.nn.rnn_cell.MultiRNNCell([cell_fn() for _ in range(params['n_layers'])])\n",
    "\n",
    "def clip_grads(loss):\n",
    "    variables = tf.trainable_variables()\n",
    "    grads = tf.gradients(loss, variables)\n",
    "    clipped_grads, _ = tf.clip_by_global_norm(grads, params['clip_norm'])\n",
    "    return zip(clipped_grads, variables)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "def forward(inputs, is_training):\n",
    "    if is_training:\n",
    "        batch_sz = tf.shape(inputs)[0]\n",
    "        \n",
    "        with tf.variable_scope('main', reuse=False):\n",
    "            embedding = tf.get_variable('lookup_table', [params['vocab_size'], params['hidden_dim']])\n",
    "            cells = multi_cell_fn()\n",
    "            \n",
    "            helper = tf.contrib.seq2seq.TrainingHelper(\n",
    "                inputs = tf.nn.embedding_lookup(embedding, inputs),\n",
    "                sequence_length = tf.count_nonzero(inputs, 1, dtype=tf.int32))\n",
    "\n",
    "            decoder = tf.contrib.seq2seq.BasicDecoder(\n",
    "                cell = cells,\n",
    "                helper = helper,\n",
    "                initial_state = cells.zero_state(batch_sz, tf.float32),\n",
    "                output_layer = tf.layers.Dense(params['vocab_size']))\n",
    "\n",
    "            decoder_output, _, _ = tf.contrib.seq2seq.dynamic_decode(\n",
    "                decoder = decoder)\n",
    "\n",
    "            logits = decoder_output.rnn_output\n",
    "            return logits\n",
    "    \n",
    "    if not is_training:\n",
    "        with tf.variable_scope('main', reuse=True):\n",
    "            cells = multi_cell_fn()\n",
    "            \n",
    "            decoder = tf.contrib.seq2seq.BeamSearchDecoder(\n",
    "                cell = cells,\n",
    "                embedding = tf.get_variable('lookup_table'),\n",
    "                start_tokens = tf.tile(tf.constant(\n",
    "                    [params['char2idx']['<start>']], dtype=tf.int32), [1]),\n",
    "                end_token = params['char2idx']['<end>'],\n",
    "                initial_state = tf.contrib.seq2seq.tile_batch(\n",
    "                    cells.zero_state(1, tf.float32), params['beam_width']),\n",
    "                beam_width = params['beam_width'],\n",
    "                output_layer = tf.layers.Dense(params['vocab_size'], _reuse=True))\n",
    "\n",
    "            decoder_out, _, _ = tf.contrib.seq2seq.dynamic_decode(\n",
    "                decoder = decoder,\n",
    "                maximum_iterations = params['seq_len'])\n",
    "\n",
    "            predict = decoder_out.predicted_ids[:, :, 0]\n",
    "            return predict"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Vocabulary size: 86\n"
     ]
    }
   ],
   "source": [
    "ints, params['char2idx'] = parse_text('../temp/anna.txt')\n",
    "params['vocab_size'] = len(params['char2idx'])\n",
    "params['idx2char'] = {i: c for c, i in params['char2idx'].items()}\n",
    "print('Vocabulary size:', params['vocab_size'])\n",
    "\n",
    "ops = {}\n",
    "X = input_fn(ints)\n",
    "\n",
    "logits = forward(start_sent(X), is_training=True)\n",
    "\n",
    "ops['global_step'] = tf.Variable(0, trainable=False)\n",
    "\n",
    "targets = end_sent(X)\n",
    "ops['loss'] = tf.reduce_mean(tf.contrib.seq2seq.sequence_loss(\n",
    "    logits = logits,\n",
    "    targets = targets,\n",
    "    weights = tf.to_float(tf.ones_like(targets))))\n",
    "\n",
    "ops['train'] = tf.train.AdamOptimizer().apply_gradients(\n",
    "    clip_grads(ops['loss']), global_step=ops['global_step'])\n",
    "\n",
    "ops['generate'] = forward(None, is_training=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Step 1 | Loss 4.450\n",
      "Step 10 | Loss 3.517\n",
      "Step 20 | Loss 3.152\n",
      "Step 30 | Loss 3.021\n",
      "Step 40 | Loss 2.928\n",
      "Step 50 | Loss 2.820\n",
      "Step 60 | Loss 2.688\n",
      "Step 70 | Loss 2.577\n",
      "Step 80 | Loss 2.487\n",
      "Step 90 | Loss 2.420\n",
      "Step 100 | Loss 2.364\n",
      "\n",
      "the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the \n",
      "\n",
      "Step 110 | Loss 2.314\n",
      "Step 120 | Loss 2.266\n",
      "Step 130 | Loss 2.215\n",
      "Step 140 | Loss 2.166\n",
      "Step 150 | Loss 2.117\n",
      "Step 160 | Loss 2.072\n",
      "Step 170 | Loss 2.027\n",
      "Step 180 | Loss 1.984\n",
      "Step 190 | Loss 1.935\n",
      "Step 200 | Loss 1.889\n",
      "\n",
      "ing and that the said, and that the said, and that the said, and that the said, and that the said, and that the said, and that the said, and that the said, and that the said, and that the said, and th\n",
      "\n",
      "Step 210 | Loss 1.842\n",
      "Step 220 | Loss 1.805\n",
      "Step 230 | Loss 1.765\n",
      "Step 240 | Loss 1.725\n",
      "Step 250 | Loss 1.683\n",
      "Step 260 | Loss 1.642\n",
      "Step 270 | Loss 1.602\n",
      "Step 280 | Loss 1.564\n",
      "Step 290 | Loss 1.523\n",
      "Step 300 | Loss 1.489\n",
      "\n",
      "ing to herself, and that he said, and that he said, and that he said, and that he said, and that he said, and that he said, and that he said, and that he said, and that he said, and that he said, and \n",
      "\n",
      "Step 310 | Loss 1.452\n",
      "Step 320 | Loss 1.411\n",
      "Step 330 | Loss 1.371\n",
      "Step 340 | Loss 1.335\n",
      "Step 350 | Loss 1.295\n",
      "Step 360 | Loss 1.259\n",
      "Step 370 | Loss 1.218\n",
      "Step 380 | Loss 1.179\n",
      "Step 390 | Loss 1.136\n",
      "Step 400 | Loss 1.101\n",
      "\n",
      "ed the children, and that he had been in the door.\n",
      "\n",
      "Stepan Arkadyevitch was conscious that he was conscious that he said, and the children, and that he had been in the door.\n",
      "\n",
      "Stepan Arkadyevitch said \n",
      "\n",
      "Step 410 | Loss 1.057\n",
      "Step 420 | Loss 1.021\n",
      "Step 430 | Loss 0.982\n",
      "Step 440 | Loss 0.962\n",
      "Step 450 | Loss 0.915\n",
      "Step 460 | Loss 0.882\n",
      "Step 470 | Loss 0.844\n",
      "Step 480 | Loss 0.822\n",
      "Step 490 | Loss 0.787\n",
      "Step 500 | Loss 0.752\n",
      "\n",
      "the could not go on like this,\" he said to himself, reconciliation with his wife, he could not go on like this,\" he said.\n",
      "\n",
      "Stepan Arkadyevitch could not get the door.\n",
      "\n",
      "Stepan Arkadyevitch could not go\n",
      "\n",
      "Step 510 | Loss 0.726\n",
      "Step 520 | Loss 0.707\n",
      "Step 530 | Loss 0.672\n",
      "Step 540 | Loss 0.645\n",
      "Step 550 | Loss 0.618\n",
      "Step 560 | Loss 0.612\n",
      "Step 570 | Loss 0.577\n",
      "Step 580 | Loss 0.555\n",
      "Step 590 | Loss 0.542\n",
      "Step 600 | Loss 0.528\n",
      "\n",
      "the children, Stepan Arkadyevitch could not get through even a short service without their father.\n",
      "\n",
      "Stepan Arkadyevitch was fond of a love him for her husband, she dropped her hands, and was impossibl\n",
      "\n",
      "Step 610 | Loss 0.508\n",
      "Step 620 | Loss 0.492\n",
      "Step 630 | Loss 0.481\n",
      "Step 640 | Loss 0.471\n",
      "Step 650 | Loss 0.455\n",
      "Step 660 | Loss 0.445\n",
      "Step 670 | Loss 0.431\n",
      "Step 680 | Loss 0.420\n",
      "Step 690 | Loss 0.409\n",
      "Step 700 | Loss 0.402\n",
      "\n",
      "the children, and for that reason I would do anything in the\n",
      "world to save them, but I don't myself know how to save them, but I don't myself know how to save them. By taking\n",
      "them away from their fath\n",
      "\n",
      "Step 710 | Loss 0.411\n",
      "Step 720 | Loss 0.415\n",
      "Step 730 | Loss 0.402\n",
      "Step 740 | Loss 0.398\n",
      "Step 750 | Loss 0.394\n",
      "Step 760 | Loss 0.387\n",
      "Step 770 | Loss 0.387\n",
      "Step 780 | Loss 0.376\n",
      "Step 790 | Loss 0.375\n",
      "Step 800 | Loss 0.381\n",
      "\n",
      "hings and from his poscaiset, and in the same a him again, and his unquestionable\n",
      "honesty. In him, in his handsome, radiant figure, his sparkling eyes,\n",
      "black hair and eyebrows, and the white and red o\n",
      "\n",
      "Step 810 | Loss 0.373\n",
      "Step 820 | Loss 0.371\n",
      "Step 830 | Loss 0.364\n",
      "Step 840 | Loss 0.363\n",
      "Step 850 | Loss 0.356\n",
      "Step 860 | Loss 0.349\n",
      "Step 870 | Loss 0.343\n",
      "Step 880 | Loss 0.346\n",
      "Step 890 | Loss 0.339\n",
      "Step 900 | Loss 0.335\n",
      "\n",
      "e she was, and what she was doing, and getting up\n",
      "rapidly, she moved towards the door.\n",
      "\n",
      "\"Well, she loves my child,\" he thought, noticin<end><end><end><end><end><end><end><end><end><end><end><end><end><end><end><end><end><end><end><end><end><end><end><end><end><end><end><end><end><end><end><end><end><end><end><end><end><end><end><end><end><end><end><end><end><end><end><end><end><end><end><end><end><end><end><end><end><end><end><end><end><end><end><end><end>\n",
      "\n",
      "Step 910 | Loss 0.343\n",
      "Step 920 | Loss 0.337\n",
      "Step 930 | Loss 0.335\n",
      "Step 940 | Loss 0.336\n",
      "Step 950 | Loss 0.349\n",
      "Step 960 | Loss 0.345\n",
      "Step 970 | Loss 0.337\n",
      "Step 980 | Loss 0.345\n",
      "Step 990 | Loss 0.339\n",
      "Step 1000 | Loss 0.342\n",
      "\n",
      "that understand that it was improper to pass judgment prematurely, and made\n",
      "him no reply.\n",
      "\n",
      "\"Who was that came in<end><end><end><end><end><end><end><end><end><end><end><end><end><end><end><end><end><end><end><end><end><end><end><end><end><end><end><end><end><end><end><end><end><end><end><end><end><end><end><end><end><end><end><end><end><end><end><end><end><end><end><end><end><end><end><end><end><end><end><end><end><end><end><end><end><end><end><end><end><end><end><end><end><end><end><end><end><end><end><end><end><end><end><end><end><end><end><end>\n",
      "\n",
      "Step 1010 | Loss 0.348\n",
      "Step 1020 | Loss 0.343\n",
      "Step 1030 | Loss 0.348\n",
      "Step 1040 | Loss 0.346\n",
      "Step 1050 | Loss 0.349\n",
      "Step 1060 | Loss 0.357\n",
      "Step 1070 | Loss 0.367\n",
      "Step 1080 | Loss 0.370\n",
      "Step 1090 | Loss 0.368\n",
      "Step 1100 | Loss 0.361\n",
      "\n",
      " and thereful for him blushed, and her, was up,\" said the distrieviv-----chollame s<end><end><end><end><end><end><end><end><end><end><end><end><end><end><end><end><end><end><end><end><end><end><end><end><end><end><end><end><end><end><end><end><end><end><end><end><end><end><end><end><end><end><end><end><end><end><end><end><end><end><end><end><end><end><end><end><end><end><end><end><end><end><end><end><end><end><end><end><end><end><end><end><end><end><end><end><end><end><end><end><end><end><end><end><end><end><end><end><end><end><end><end><end><end><end><end><end><end><end><end><end><end><end><end><end><end><end><end><end><end><end><end><end><end><end><end><end>\n",
      "\n",
      "Step 1110 | Loss 0.354\n",
      "Step 1120 | Loss 0.360\n",
      "Step 1130 | Loss 0.357\n",
      "Step 1140 | Loss 0.356\n",
      "Step 1150 | Loss 0.354\n",
      "Step 1160 | Loss 0.363\n",
      "Step 1170 | Loss 0.371\n",
      "Step 1180 | Loss 0.361\n",
      "Step 1190 | Loss 0.373\n",
      "Step 1200 | Loss 0.377\n",
      "\n",
      "the winders of all the mastive to the liberalism of the district you came in.\n",
      "\n",
      "\"Yes, sir.\"\n",
      "\n",
      "Stepan Arkadyevitch put on his fur coat and went out onto the steps.\n",
      "\n",
      "\"You won't dine at home?\" said Matvey,\n",
      "\n",
      "Step 1210 | Loss 0.376\n"
     ]
    }
   ],
   "source": [
    "sess = tf.Session()\n",
    "sess.run(tf.global_variables_initializer())\n",
    "while True:\n",
    "    try:\n",
    "        _, step, loss = sess.run([ops['train'], ops['global_step'], ops['loss']])\n",
    "    except tf.errors.OutOfRangeError:\n",
    "        break\n",
    "    else:\n",
    "        if step % params['display_step'] == 0 or step == 1:\n",
    "            print(\"Step %d | Loss %.3f\" % (step, loss))\n",
    "        if step % params['generate_step'] == 0 and step > 1:\n",
    "            ints = sess.run(ops['generate'])[0]\n",
    "            print('\\n'+''.join([params['idx2char'][i] for i in ints])+'\\n')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
